Generalize strided store pattern in interleave access pass

[asbirlea]

Summary

This patch aims to generalize matching of the strided store accesses to more general masks.
The more general rule is to have consecutive accesses based on the stride:
[x, y, ... z, x+1, y+1, ...z+1, x+2, y+2, ...z+2, ...]
and for each start element in each stride (x, y, ... z] to be aligned.
However all elements in the masks need not form a contiguous space, there may be gaps.
As before, undefs are allowed and filled in with adjacent element loads.

Note this patch is not final, but I would like to get feedback on the approach.
There are at least the pending TODOs.

[mssimpso]

Hi Alina,

I think I understand this, but I just want to be sure I get how this differs from what we currently have before going further. Currently, we only match [x, y, ..., z, x+1, y+1, z+1, ...] where each y-x and each z-y equals the number of sub elements for the given factor. Or said another way, if I create a list or all the x's followed by all the y's and then all the z's, the entire list would be consecutive. With your path, the only requirement is that each sub-list be consecutive. Is this right?

The current approach was designed to match the shuffle patterns produced by the loop vectorizer. I'm curious to know where we are generating these more general patterns. Have you run across some code examples?

Also, another high level comment before I start looking at the details: you'll want to include some IR test cases as well (to be run with opt instead of llc).

Matt.

[asbirlea]

Hi Matt,

Thanks for looking to review this. Please find my answers below.

In D23646#521091, @mssimpso wrote:

Hi Alina,

I think I understand this, but I just want to be sure I get how this differs from what we currently have before going further. Currently, we only match [x, y, ..., z, x+1, y+1, z+1, ...] where each y-x and each z-y equals the number of sub elements for the given factor. Or said another way, if I create a list or all the x's followed by all the y's and then all the z's, the entire list would be consecutive. With your path, the only requirement is that each sub-list be consecutive. Is this right?

That's right. Also, from my understanding, x is always 0. So all elements form a consecutive sublist which always starts at 0.
My first approach was actually to generalize this just to add a prefix to remove the "starts with 0" restriction and a more general stride that allowed gaps. But this still didn't cover all the testcases I came across, such as the example I added in "store_general_mask_factor4".

To answer your question below, the usecases I'm looking at are generated by Halide (https://github.com/halide/Halide).
Halide generates LLVM IR and relies on its optimization pipeline and lowering, but they need to generate explicit intrinsics (including strided loads and stores) for arm and aarch64, because their patterns are not lowered to intrinsics by LLVM.
Since this approach was taken before the interleaved-access pass was added, it's quite understandable, but LLVM is more powerful now and I'm trying to make use of this, and in the process, cover the cases missing in LLVM.
For example, for strided loads the interleaved-access pass does cover the code patterns generated by Halide, so the "custom" intrinsic code generation in Halide will soon be removed. My goal is to improve the pass to make this happen for the stores as well.
The tests I will add are actually simplified versions of what Halide is generating.

The current approach was designed to match the shuffle patterns produced by the loop vectorizer. I'm curious to know where we are generating these more general patterns. Have you run across some code examples?

Also, another high level comment before I start looking at the details: you'll want to include some IR test cases as well (to be run with opt instead of llc).

Agreed, the plan is to add more tests, including IR tests.

Matt.

[mssimpso]

lib/CodeGen/InterleavedAccessPass.cpp

159–164 ↗

(On Diff #68484)

You should probably update this to define the more general pattern.

181–207 ↗

(On Diff #68484)

This looks fairly reasonable to me, but the parts dealing with undef are pretty difficult to follow. I think some more high-level comments would help people better understand what's going on here.

[asbirlea]

Address comments re. comments.
Complete TODOs.
Requesting help on whether to include TLI.misalignedAccess check and on what's the correct way to do it.

[asbirlea]

lib/CodeGen/InterleavedAccessPass.cpp

400 ↗

(On Diff #70479)

This is a part that I'm not sure is needed, and how to address it.

The goal was to check for the alignment of each of the strides, i.e. BaseStoreAddress + StartingIncrementInStride, for all stride [0, Factor).
The commented attempt has a series of problems and does not achieve this. Should this check exist and what's the correct way to handle it?

[mssimpso]

Hi Alina,

Sorry for the delay. I'm not quite sure I understand this patch anymore. I'm adding Tim Northover (ARM/AArch64 code owner) as a reviewer to hopefully get this unstuck.

[asbirlea]

SG, thank you!

I'm going back to look at the alignment check this afternoon (that's the big commented out block).
I'd really like to understand why some basic alignment checks lead to ARM tests failing and not their AArch64 counterparts.

[mssimpso]

lib/CodeGen/InterleavedAccessPass.cpp

400 ↗

(On Diff #70479)

I could be wrong about this, but I don't think you need to worry about alignment here. I'm not seeing how the memory behavior with this patch would be different than the current situation.

[asbirlea]

lib/CodeGen/InterleavedAccessPass.cpp

400 ↗

(On Diff #70479)

I agree with you that there should be no significant difference from the current situation. There is one small difference though...before there could be one misaligned access, now there may be Factor such accesses.
That's why I'd still like to understand the alignment issue - whether the check is needed or not and in what form. Perhaps it would be better to have it in a separate patch though.

[asbirlea]

Remove comment block checking for alignment. Will revisit in a future patch.

[asbirlea]

Pinging patch.

Also, working around the case when masks are larger than 16 elements.
This can happen now if Halide takes advantage of this pass for strided stores.
This is not the only use-case of larger shuffle masks, but the topic is beyond the scope of this patch.

[asbirlea]

Minor edit of temporary variables.

[rengolin]

lib/CodeGen/InterleavedAccessPass.cpp

193 ↗

(On Diff #73332)

Nit, ij is hard to follow. Try lane or something more expressive. (this is not a matrix :)

204 ↗

(On Diff #73332)

This is really confusing. Can you factor the comparison elements out with expressive names, so the if becomes a comparison of obvious terms?

208 ↗

(On Diff #73332)

PreviousMask is always used in conjunction with PreviousPos, so you don't need the mask to be signed and you can compare the pos in the block above and get rid of the static casts. Or you could have an additional boolean flag and make them both unsigned.

lib/Target/AArch64/AArch64ISelLowering.cpp

7243 ↗

(On Diff #73332)

Nit. use brackets here:

if (...) {
  ...
} else {

7247 ↗

(On Diff #73332)

I don't get the - j here.

lib/Target/ARM/ARMISelLowering.cpp

13126 ↗

(On Diff #73332)

Better to duplicate the comment, I think. These back-ends evolve at different paces.

test/CodeGen/ARM/arm-interleaved-accesses.ll

321 ↗

(On Diff #73332)

is this really guaranteed to reproduce? they don't seem connected to the pcs directly...

[asbirlea]

Address comments.

[asbirlea]

lib/CodeGen/InterleavedAccessPass.cpp

193 ↗

(On Diff #73332)

Renamed to Lane. I hope changing the NumSubElts to LaneLen makes more sense too.
I'm inclined to change it in the lowering files as well for consistency.

lib/Target/AArch64/AArch64ISelLowering.cpp

7247 ↗

(On Diff #73332)

Assuming the mask starts with a few undefs, this computes what the start of the mask would be based on the first non-undef value.
The computation is done first in the pass to make sure the start is a positive value (hence the correctness comment below on "StartMask cannot be negative")

test/CodeGen/ARM/arm-interleaved-accesses.ll

321 ↗

(On Diff #73332)

I'm not sure about this TBH, and not sure how to verify it.
Should I replace it by a simple vst4.32 check?

[rengolin]

Hi Alina,

This is looking much better, thanks!

The code has a lot of undef handling, but not much in the way of testing it. I think we should have at least the following:

• one and two undefs in the middle
• one undef at the beginning and one at the end
• all undefs in one lane
• one undef in each lane, at different positions

Repeating the pattern of your current tests but adding undefs should be enough.

cheers,
--renato

lib/CodeGen/InterleavedAccessPass.cpp

193 ↗

(On Diff #73332)

Nice, much better! I agree with renaming the lowering code, too.

188 ↗

(On Diff #74721)

Better to declare I and J inside the for declaration.

200 ↗

(On Diff #74721)

A comment here would help...

// If both defined, only sequential values allowed

216 ↗

(On Diff #74721)

What about the case where the first lane is undef, but the others aren't?

226 ↗

(On Diff #74721)

Instead of using a SavedNonUndef above, you could save the last non-undef value and the number of undefs since that value. That'd make the next-value computation easier:

If (NextValue != SavedValue + NumUndefs)
  break;

and also help get the StartMask here, for free.

237 ↗

(On Diff #74721)

"Found an interleaved..."

lib/Target/AArch64/AArch64ISelLowering.cpp

7247 ↗

(On Diff #73332)

Right, I agree you could repeat the naming pattern above, here.

test/CodeGen/ARM/arm-interleaved-accesses.ll

321 ↗

(On Diff #73332)

Something like:

vst4.32 {d{{\n+}}, d{{\n+}}, d{{\n+}}, d{{\n+}}}, [r0]

would do. (I'm not sure of the triple brackets there...)

[asbirlea]

Address comments. One pending.

[asbirlea]

Great point on the lack of testing. But I wasn't happy with the coverage the vst4/st4 had.
I added a pattern for vst3/st3 that covers the undefs in the middle of a lane.

lib/CodeGen/InterleavedAccessPass.cpp

188 ↗

(On Diff #74721)

There's a check on I and J following each loop. I could add an additional flag to check that we broke out of the loop early, but it seemed overkill to do that when I and J could be used if declared outside the loop.

216 ↗

(On Diff #74721)

Nothing wrong with that (unless I'm missing something)..
It'll check the correctness for the ones that follow and the first one will receive a value based on the following values - that's the start mask value.

226 ↗

(On Diff #74721)

I'm still looking into this one.
I can do without SaveNonUndef, and update the condition to a "SavedLaneValue+SavedNoUndefs (+1)".
This needs an additional if clause in the loop to increment the SavedNoUndefs, and at least another check to help with computing the mask. The second check is because right now I only store SavedLaneValue if a value is followed by an undef, but at the end of the mask we'll need this updated too to get the correct StartMask as something like SavedLaneValue+SavedNoUndefs -LaneLen (+/- 1).
Right now I find it easier to just compute the StartMask in the same j loop.
So, yeah, still looking what's the cleanest way to do this.

[asbirlea]

Address remaining comment. Add additional testcases.

[asbirlea]

[clang-format]

[asbirlea]

Gentle ping.

[asbirlea]

Re-pinging patch.

[asbirlea]

Pinging again.

[rengolin]

Hi,

Sorry to keep you waiting, this completely fell out of my radar.

I think the code looks good now, just need to make sure the test is generic enough on the CHECK line (see inline comment).

cheers,
--renato

test/CodeGen/AArch64/aarch64-interleaved-accesses.ll

285 ↗

(On Diff #75332)

Please, also use:

st4.32 {d{{\n+}}, d{{\n+}}, d{{\n+}}, d{{\n+}}}, [x0]

here.

[asbirlea]

Update aarch64 test.

[asbirlea]

Thank you for the review, Renato!

llvm/trunk/lib/CodeGen/InterleavedAccessPass.cpp

static bool isDeInterleaveMask(ArrayRef<int> Mask, unsigned &Factor,
   // Check potential Factors.
   for (Factor = 2; Factor <= MaxFactor; Factor++)
     if (isDeInterleaveMaskOfFactor(Mask, Factor, Index))
       return true;
 
   return false;
 }
 
-/// \brief Check if the mask is RE-interleave mask for an interleaved store.
-///
-/// I.e. <0, NumSubElts, ... , NumSubElts*(Factor - 1), 1, NumSubElts + 1, ...>
+/// \brief Check if the mask can be used in an interleaved store.
+//
+/// It checks for a more general pattern than the RE-interleave mask.
+/// I.e. <x, y, ... z, x+1, y+1, ...z+1, x+2, y+2, ...z+2, ...>
+/// E.g. For a Factor of 2 (LaneLen=4): <4, 32, 5, 33, 6, 34, 7, 35>
+/// E.g. For a Factor of 3 (LaneLen=4): <4, 32, 16, 5, 33, 17, 6, 34, 18, 7, 35, 19>
+/// E.g. For a Factor of 4 (LaneLen=2): <8, 2, 12, 4, 9, 3, 13, 5>
 ///
-/// E.g. The RE-interleave mask (Factor = 2) could be:
-///     <0, 4, 1, 5, 2, 6, 3, 7>
+/// The particular case of an RE-interleave mask is:
+/// I.e. <0, LaneLen, ... , LaneLen*(Factor - 1), 1, LaneLen + 1, ...>
+/// E.g. For a Factor of 2 (LaneLen=4): <0, 4, 1, 5, 2, 6, 3, 7>
 static bool isReInterleaveMask(ArrayRef<int> Mask, unsigned &Factor,
                                unsigned MaxFactor) {
   unsigned NumElts = Mask.size();
   if (NumElts < 4)
     return false;
 
   // Check potential Factors.
   for (Factor = 2; Factor <= MaxFactor; Factor++) {
     if (NumElts % Factor)
       continue;
 
-    unsigned NumSubElts = NumElts / Factor;
-    if (!isPowerOf2_32(NumSubElts))
+    unsigned LaneLen = NumElts / Factor;
+    if (!isPowerOf2_32(LaneLen))
       continue;
 
-    // Check whether each element matchs the RE-interleaved rule. Ignore undef
-    // elements.
-    unsigned i = 0;
-    for (; i < NumElts; i++)
-      if (Mask[i] >= 0 &&
-          static_cast<unsigned>(Mask[i]) !=
-              (i % Factor) * NumSubElts + i / Factor)
+    // Check whether each element matches the general interleaved rule.
+    // Ignore undef elements, as long as the defined elements match the rule.
+    // Outer loop processes all factors (x, y, z in the above example)
+    unsigned I = 0, J;
+    for (; I < Factor; I++) {
+      unsigned SavedLaneValue;
+      unsigned SavedNoUndefs = 0;
+
+      // Inner loop processes consecutive accesses (x, x+1... in the example)
+      for (J = 0; J < LaneLen - 1; J++) {
+        // Lane computes x's position in the Mask
+        unsigned Lane = J * Factor + I;
+        unsigned NextLane = Lane + Factor;
+        int LaneValue = Mask[Lane];
+        int NextLaneValue = Mask[NextLane];
+
+        // If both are defined, values must be sequential
+        if (LaneValue >= 0 && NextLaneValue >= 0 &&
+            LaneValue + 1 != NextLaneValue)
           break;
 
-    // Find a RE-interleaved mask of current factor.
-    if (i == NumElts)
+        // If the next value is undef, save the current one as reference
+        if (LaneValue >= 0 && NextLaneValue < 0) {
+          SavedLaneValue = LaneValue;
+          SavedNoUndefs = 1;
+        }
+
+        // Undefs are allowed, but defined elements must still be consecutive:
+        // i.e.: x,..., undef,..., x + 2,..., undef,..., undef,..., x + 5, ....
+        // Verify this by storing the last non-undef followed by an undef
+        // Check that following non-undef masks are incremented with the
+        // corresponding distance.
+        if (SavedNoUndefs > 0 && LaneValue < 0) {
+          SavedNoUndefs++;
+          if (NextLaneValue >= 0 &&
+              SavedLaneValue + SavedNoUndefs != (unsigned)NextLaneValue)
+            break;
+        }
+      }
+
+      if (J < LaneLen - 1)
+        break;
+
+      int StartMask = 0;
+      if (Mask[I] >= 0) {
+        // Check that the start of the I range (J=0) is greater than 0
+        StartMask = Mask[I];
+      } else if (Mask[(LaneLen - 1) * Factor + I] >= 0) {
+        // StartMask defined by the last value in lane
+        StartMask = Mask[(LaneLen - 1) * Factor + I] - J;
+      } else if (SavedNoUndefs > 0) {
+        // StartMask defined by some non-zero value in the j loop
+        StartMask = SavedLaneValue - (LaneLen - 1 - SavedNoUndefs);
+      }
+      // else StartMask remains set to 0, i.e. all elements are undefs
+
+      if (StartMask < 0)
+        break;
+    }
+
+    // Found an interleaved mask of current factor.
+    if (I == Factor)
       return true;
   }
 
   return false;
 }
 
 bool InterleavedAccess::lowerInterleavedLoad(
     LoadInst *LI, SmallVector<Instruction *, 32> &DeadInsts) {

llvm/trunk/lib/Target/AArch64/AArch64ISelLowering.cpp

static Constant *getSequentialMask(IRBuilder<> &Builder, unsigned Start,
 
   return ConstantVector::get(Mask);
 }
 
 /// \brief Lower an interleaved store into a stN intrinsic.
 ///
 /// E.g. Lower an interleaved store (Factor = 3):
 ///        %i.vec = shuffle <8 x i32> %v0, <8 x i32> %v1,
-///                                  <0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11>
+///                 <0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11>
 ///        store <12 x i32> %i.vec, <12 x i32>* %ptr
 ///
 ///      Into:
 ///        %sub.v0 = shuffle <8 x i32> %v0, <8 x i32> v1, <0, 1, 2, 3>
 ///        %sub.v1 = shuffle <8 x i32> %v0, <8 x i32> v1, <4, 5, 6, 7>
 ///        %sub.v2 = shuffle <8 x i32> %v0, <8 x i32> v1, <8, 9, 10, 11>
 ///        call void llvm.aarch64.neon.st3(%sub.v0, %sub.v1, %sub.v2, %ptr)
 ///
 /// Note that the new shufflevectors will be removed and we'll only generate one
 /// st3 instruction in CodeGen.
+///
+/// Example for a more general valid mask (Factor 3). Lower:
+///        %i.vec = shuffle <32 x i32> %v0, <32 x i32> %v1,
+///                 <4, 32, 16, 5, 33, 17, 6, 34, 18, 7, 35, 19>
+///        store <12 x i32> %i.vec, <12 x i32>* %ptr
+///
+///      Into:
+///        %sub.v0 = shuffle <32 x i32> %v0, <32 x i32> v1, <4, 5, 6, 7>
+///        %sub.v1 = shuffle <32 x i32> %v0, <32 x i32> v1, <32, 33, 34, 35>
+///        %sub.v2 = shuffle <32 x i32> %v0, <32 x i32> v1, <16, 17, 18, 19>
+///        call void llvm.aarch64.neon.st3(%sub.v0, %sub.v1, %sub.v2, %ptr)
 bool AArch64TargetLowering::lowerInterleavedStore(StoreInst *SI,
                                                   ShuffleVectorInst *SVI,
                                                   unsigned Factor) const {
   assert(Factor >= 2 && Factor <= getMaxSupportedInterleaveFactor() &&
          "Invalid interleave factor");
 
   VectorType *VecTy = SVI->getType();
   assert(VecTy->getVectorNumElements() % Factor == 0 &&
          "Invalid interleaved store");
 
-  unsigned NumSubElts = VecTy->getVectorNumElements() / Factor;
+  unsigned LaneLen = VecTy->getVectorNumElements() / Factor;
   Type *EltTy = VecTy->getVectorElementType();
-  VectorType *SubVecTy = VectorType::get(EltTy, NumSubElts);
+  VectorType *SubVecTy = VectorType::get(EltTy, LaneLen);
 
   const DataLayout &DL = SI->getModule()->getDataLayout();
   unsigned SubVecSize = DL.getTypeSizeInBits(SubVecTy);
 
   // Skip if we do not have NEON and skip illegal vector types.
   if (!Subtarget->hasNEON() || (SubVecSize != 64 && SubVecSize != 128))
     return false;
 
   Value *Op0 = SVI->getOperand(0);
   Value *Op1 = SVI->getOperand(1);
   IRBuilder<> Builder(SI);
 
   // StN intrinsics don't support pointer vectors as arguments. Convert pointer
   // vectors to integer vectors.
   if (EltTy->isPointerTy()) {
     Type *IntTy = DL.getIntPtrType(EltTy);
     unsigned NumOpElts =
         dyn_cast<VectorType>(Op0->getType())->getVectorNumElements();
 
     // Convert to the corresponding integer vector.
     Type *IntVecTy = VectorType::get(IntTy, NumOpElts);
     Op0 = Builder.CreatePtrToInt(Op0, IntVecTy);
     Op1 = Builder.CreatePtrToInt(Op1, IntVecTy);
 
-    SubVecTy = VectorType::get(IntTy, NumSubElts);
+    SubVecTy = VectorType::get(IntTy, LaneLen);
   }
 
   Type *PtrTy = SubVecTy->getPointerTo(SI->getPointerAddressSpace());
   Type *Tys[2] = {SubVecTy, PtrTy};
   static const Intrinsic::ID StoreInts[3] = {Intrinsic::aarch64_neon_st2,
                                              Intrinsic::aarch64_neon_st3,
                                              Intrinsic::aarch64_neon_st4};
   Function *StNFunc =
       Intrinsic::getDeclaration(SI->getModule(), StoreInts[Factor - 2], Tys);
 
   SmallVector<Value *, 5> Ops;
 
   // Split the shufflevector operands into sub vectors for the new stN call.
-  for (unsigned i = 0; i < Factor; i++)
+  auto Mask = SVI->getShuffleMask();
+  for (unsigned i = 0; i < Factor; i++) {
+    if (Mask[i] >= 0) {
+      Ops.push_back(Builder.CreateShuffleVector(
+        Op0, Op1, getSequentialMask(Builder, Mask[i], LaneLen)));
+    } else {
+      unsigned StartMask = 0;
+      for (unsigned j = 1; j < LaneLen; j++) {
+        if (Mask[j*Factor + i] >= 0) {
+          StartMask = Mask[j*Factor + i] - j;
+          break;
+        }
+      }
+      // Note: If all elements in a chunk are undefs, StartMask=0!
+      // Note: Filling undef gaps with random elements is ok, since
+      // those elements were being written anyway (with undefs).
+      // In the case of all undefs we're defaulting to using elems from 0
+      // Note: StartMask cannot be negative, it's checked in isReInterleaveMask
       Ops.push_back(Builder.CreateShuffleVector(
-        Op0, Op1, getSequentialMask(Builder, NumSubElts * i, NumSubElts)));
+        Op0, Op1, getSequentialMask(Builder, StartMask, LaneLen)));
+    }
+  }
 
   Ops.push_back(Builder.CreateBitCast(SI->getPointerOperand(), PtrTy));
   Builder.CreateCall(StNFunc, Ops);
   return true;
 }
 
 static bool memOpAlign(unsigned DstAlign, unsigned SrcAlign,
                        unsigned AlignCheck) {

llvm/trunk/lib/Target/ARM/ARMISelLowering.cpp

 ///      Into:
 ///        %sub.v0 = shuffle <8 x i32> %v0, <8 x i32> v1, <0, 1, 2, 3>
 ///        %sub.v1 = shuffle <8 x i32> %v0, <8 x i32> v1, <4, 5, 6, 7>
 ///        %sub.v2 = shuffle <8 x i32> %v0, <8 x i32> v1, <8, 9, 10, 11>
 ///        call void llvm.arm.neon.vst3(%ptr, %sub.v0, %sub.v1, %sub.v2, 4)
 ///
 /// Note that the new shufflevectors will be removed and we'll only generate one
 /// vst3 instruction in CodeGen.
+///
+/// Example for a more general valid mask (Factor 3). Lower:
+///        %i.vec = shuffle <32 x i32> %v0, <32 x i32> %v1,
+///                 <4, 32, 16, 5, 33, 17, 6, 34, 18, 7, 35, 19>
+///        store <12 x i32> %i.vec, <12 x i32>* %ptr
+///
+///      Into:
+///        %sub.v0 = shuffle <32 x i32> %v0, <32 x i32> v1, <4, 5, 6, 7>
+///        %sub.v1 = shuffle <32 x i32> %v0, <32 x i32> v1, <32, 33, 34, 35>
+///        %sub.v2 = shuffle <32 x i32> %v0, <32 x i32> v1, <16, 17, 18, 19>
+///        call void llvm.arm.neon.vst3(%ptr, %sub.v0, %sub.v1, %sub.v2, 4)
 bool ARMTargetLowering::lowerInterleavedStore(StoreInst *SI,
                                               ShuffleVectorInst *SVI,
                                               unsigned Factor) const {
   assert(Factor >= 2 && Factor <= getMaxSupportedInterleaveFactor() &&
          "Invalid interleave factor");
 
   VectorType *VecTy = SVI->getType();
   assert(VecTy->getVectorNumElements() % Factor == 0 &&
          "Invalid interleaved store");
 
-  unsigned NumSubElts = VecTy->getVectorNumElements() / Factor;
+  unsigned LaneLen = VecTy->getVectorNumElements() / Factor;
   Type *EltTy = VecTy->getVectorElementType();
-  VectorType *SubVecTy = VectorType::get(EltTy, NumSubElts);
+  VectorType *SubVecTy = VectorType::get(EltTy, LaneLen);
 
   const DataLayout &DL = SI->getModule()->getDataLayout();
   unsigned SubVecSize = DL.getTypeSizeInBits(SubVecTy);
   bool EltIs64Bits = DL.getTypeSizeInBits(EltTy) == 64;
 
   // Skip if we do not have NEON and skip illegal vector types and vector types
   // with i64/f64 elements (vstN doesn't support i64/f64 elements).
   if (!Subtarget->hasNEON() || (SubVecSize != 64 && SubVecSize != 128) ||
if (EltTy->isPointerTy()) {
     Type *IntTy = DL.getIntPtrType(EltTy);
 
     // Convert to the corresponding integer vector.
     Type *IntVecTy =
         VectorType::get(IntTy, Op0->getType()->getVectorNumElements());
     Op0 = Builder.CreatePtrToInt(Op0, IntVecTy);
     Op1 = Builder.CreatePtrToInt(Op1, IntVecTy);
 
-    SubVecTy = VectorType::get(IntTy, NumSubElts);
+    SubVecTy = VectorType::get(IntTy, LaneLen);
   }
 
   static const Intrinsic::ID StoreInts[3] = {Intrinsic::arm_neon_vst2,
                                              Intrinsic::arm_neon_vst3,
                                              Intrinsic::arm_neon_vst4};
   SmallVector<Value *, 6> Ops;
 
   Type *Int8Ptr = Builder.getInt8PtrTy(SI->getPointerAddressSpace());
   Ops.push_back(Builder.CreateBitCast(SI->getPointerOperand(), Int8Ptr));
 
   Type *Tys[] = { Int8Ptr, SubVecTy };
   Function *VstNFunc = Intrinsic::getDeclaration(
       SI->getModule(), StoreInts[Factor - 2], Tys);
 
   // Split the shufflevector operands into sub vectors for the new vstN call.
-  for (unsigned i = 0; i < Factor; i++)
+  auto Mask = SVI->getShuffleMask();
+  for (unsigned i = 0; i < Factor; i++) {
+    if (Mask[i] >= 0) {
       Ops.push_back(Builder.CreateShuffleVector(
-        Op0, Op1, getSequentialMask(Builder, NumSubElts * i, NumSubElts)));
+        Op0, Op1, getSequentialMask(Builder, Mask[i], LaneLen)));
+    } else {
+      unsigned StartMask = 0;
+      for (unsigned j = 1; j < LaneLen; j++) {
+        if (Mask[j*Factor + i] >= 0) {
+          StartMask = Mask[j*Factor + i] - j;
+          break;
+        }
+      }
+      // Note: If all elements in a chunk are undefs, StartMask=0!
+      // Note: Filling undef gaps with random elements is ok, since
+      // those elements were being written anyway (with undefs).
+      // In the case of all undefs we're defaulting to using elems from 0
+      // Note: StartMask cannot be negative, it's checked in isReInterleaveMask
+      Ops.push_back(Builder.CreateShuffleVector(
+        Op0, Op1, getSequentialMask(Builder, StartMask, LaneLen)));
+    }
+  }
 
   Ops.push_back(Builder.getInt32(SI->getAlignment()));
   Builder.CreateCall(VstNFunc, Ops);
   return true;
 }
 
 enum HABaseType {
   HA_UNKNOWN = 0,

llvm/trunk/test/CodeGen/AArch64/aarch64-interleaved-accesses.ll

 ; NONEON-LABEL: load_factor2_with_extract_user:
 ; NONEON-NOT: ld2
 define i32 @load_factor2_with_extract_user(<8 x i32>* %a) {
   %1 = load <8 x i32>, <8 x i32>* %a, align 8
   %2 = shufflevector <8 x i32> %1, <8 x i32> undef, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
   %3 = extractelement <8 x i32> %1, i32 2
   ret i32 %3
 }
+
+; NEON-LABEL: store_general_mask_factor4:
+; NEON: st4 { v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s }, [x0]
+; NONEON-LABEL: store_general_mask_factor4:
+; NONEON-NOT: st4
+define void @store_general_mask_factor4(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 4, i32 16, i32 32, i32 8, i32 5, i32 17, i32 33, i32 9>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor4_undefbeg:
+; NEON: st4 { v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s }, [x0]
+; NONEON-LABEL: store_general_mask_factor4_undefbeg:
+; NONEON-NOT: st4
+define void @store_general_mask_factor4_undefbeg(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 undef, i32 16, i32 32, i32 8, i32 5, i32 17, i32 33, i32 9>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor4_undefend:
+; NEON: st4 { v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s }, [x0]
+; NONEON-LABEL: store_general_mask_factor4_undefend:
+; NONEON-NOT: st4
+define void @store_general_mask_factor4_undefend(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 4, i32 16, i32 32, i32 8, i32 5, i32 17, i32 33, i32 undef>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor4_undefmid:
+; NEON: st4 { v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s }, [x0]
+; NONEON-LABEL: store_general_mask_factor4_undefmid:
+; NONEON-NOT: st4
+define void @store_general_mask_factor4_undefmid(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 4, i32 undef, i32 32, i32 8, i32 5, i32 17, i32 undef, i32 9>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor4_undefmulti:
+; NEON: st4 { v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s }, [x0]
+; NONEON-LABEL: store_general_mask_factor4_undefmulti:
+; NONEON-NOT: st4
+define void @store_general_mask_factor4_undefmulti(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 4, i32 undef, i32 undef, i32 8, i32 undef, i32 undef, i32 undef, i32 9>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3:
+; NEON: st3 { v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s }, [x0]
+; NONEON-LABEL: store_general_mask_factor3:
+; NONEON-NOT: st3
+define void @store_general_mask_factor3(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 4, i32 32, i32 16, i32 5, i32 33, i32 17, i32 6, i32 34, i32 18, i32 7, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_undefmultimid:
+; NEON: st3 { v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s }, [x0]
+; NONEON-LABEL: store_general_mask_factor3_undefmultimid:
+; NONEON-NOT: st3
+define void @store_general_mask_factor3_undefmultimid(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 4, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 7, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_undef_fail:
+; NEON-NOT: st3
+; NONEON-LABEL: store_general_mask_factor3_undef_fail:
+; NONEON-NOT: st3
+define void @store_general_mask_factor3_undef_fail(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 4, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 8, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_undeflane:
+; NEON: st3 { v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s, v{{[0-9]+}}.{{[0-9]+}}s }, [x0]
+; NONEON-LABEL: store_general_mask_factor3_undeflane:
+; NONEON-NOT: st3
+define void @store_general_mask_factor3_undeflane(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 undef, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 undef, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_negativestart:
+; NEON-NOT: st3
+; NONEON-LABEL: store_general_mask_factor3_negativestart:
+; NONEON-NOT: st3
+define void @store_general_mask_factor3_negativestart(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 undef, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 2, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+

llvm/trunk/test/CodeGen/ARM/arm-interleaved-accesses.ll

 ; NONEON-LABEL: load_factor2_with_extract_user:
 ; NONEON-NOT: vld2
 define i32 @load_factor2_with_extract_user(<8 x i32>* %a) {
   %1 = load <8 x i32>, <8 x i32>* %a, align 8
   %2 = shufflevector <8 x i32> %1, <8 x i32> undef, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
   %3 = extractelement <8 x i32> %1, i32 2
   ret i32 %3
 }
+
+; NEON-LABEL: store_general_mask_factor4:
+; NEON: vst4.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor4:
+; NONEON-NOT: vst4.32
+define void @store_general_mask_factor4(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 4, i32 16, i32 32, i32 8, i32 5, i32 17, i32 33, i32 9>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor4_undefbeg:
+; NEON: vst4.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor4_undefbeg:
+; NONEON-NOT: vst4.32
+define void @store_general_mask_factor4_undefbeg(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 undef, i32 16, i32 32, i32 8, i32 5, i32 17, i32 33, i32 9>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor4_undefend:
+; NEON: vst4.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor4_undefend:
+; NONEON-NOT: vst4.32
+define void @store_general_mask_factor4_undefend(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 4, i32 16, i32 32, i32 8, i32 5, i32 17, i32 33, i32 undef>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor4_undefmid:
+; NEON: vst4.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor4_undefmid:
+; NONEON-NOT: vst4.32
+define void @store_general_mask_factor4_undefmid(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 4, i32 undef, i32 32, i32 8, i32 5, i32 17, i32 undef, i32 9>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor4_undefmulti:
+; NEON: vst4.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor4_undefmulti:
+; NONEON-NOT: vst4.32
+define void @store_general_mask_factor4_undefmulti(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <8 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <8 x i32> <i32 4, i32 undef, i32 undef, i32 8, i32 undef, i32 undef, i32 undef, i32 9>
+  store <8 x i32> %i.vec, <8 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3:
+; NEON: vst3.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor3:
+; NONEON-NOT: vst3.32
+define void @store_general_mask_factor3(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 4, i32 32, i32 16, i32 5, i32 33, i32 17, i32 6, i32 34, i32 18, i32 7, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_undefmultimid:
+; NEON: vst3.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor3_undefmultimid:
+; NONEON-NOT: vst3.32
+define void @store_general_mask_factor3_undefmultimid(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 4, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 7, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_undef_fail:
+; NEON-NOT: vst3.32
+; NONEON-LABEL: store_general_mask_factor3_undef_fail:
+; NONEON-NOT: vst3.32
+define void @store_general_mask_factor3_undef_fail(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 4, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 8, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_undeflane:
+; NEON: vst3.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor3_undeflane:
+; NONEON-NOT: vst3.32
+define void @store_general_mask_factor3_undeflane(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 undef, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 undef, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_endstart_fail:
+; NEON-NOT: vst3.32
+; NONEON-LABEL: store_general_mask_factor3_endstart_fail:
+; NONEON-NOT: vst3.32
+define void @store_general_mask_factor3_endstart_fail(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 undef, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 2, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_endstart_pass:
+; NEON: vst3.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor3_endstart_pass:
+; NONEON-NOT: vst3.32
+define void @store_general_mask_factor3_endstart_pass(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 undef, i32 32, i32 16, i32 undef, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 7, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_midstart_fail:
+; NEON-NOT: vst3.32
+; NONEON-LABEL: store_general_mask_factor3_midstart_fail:
+; NONEON-NOT: vst3.32
+define void @store_general_mask_factor3_midstart_fail(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 undef, i32 32, i32 16, i32 0, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 undef, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+
+; NEON-LABEL: store_general_mask_factor3_midstart_pass:
+; NEON: vst3.32 {d{{[0-9]+}}, d{{[0-9]+}}, d{{[0-9]+}}}, [r0]
+; NONEON-LABEL: store_general_mask_factor3_midstart_pass:
+; NONEON-NOT: vst3.32
+define void @store_general_mask_factor3_midstart_pass(i32* %ptr, <32 x i32> %v0, <32 x i32> %v1) {
+  %base = bitcast i32* %ptr to <12 x i32>*
+  %i.vec = shufflevector <32 x i32> %v0, <32 x i32> %v1, <12 x i32> <i32 undef, i32 32, i32 16, i32 1, i32 33, i32 17, i32 undef, i32 34, i32 18, i32 undef, i32 35, i32 19>
+  store <12 x i32> %i.vec, <12 x i32>* %base, align 4
+  ret void
+}
+