[AArch64] Add index check before lowerInterleavedStore() uses ShuffleVectorInst's mask

[Peter]

Summary

This commit fixes https://github.com/llvm/llvm-project/issues/57326.

Currently we would take a Mask out and directly use it by doing auto Mask = SVI->getShuffleMask();
However, if the mask is undef, this Mask is not initialized. It might be a vector of -1 or random integers.
This would cause an Out-of-bound read laterwhen trying to find a StartMask.

This change checks if all indices in the Mask is in the allowed range. If not, we would early abort.

[fhahn]

llvm/test/CodeGen/AArch64/aarch64-shuffle-undef.ll

24 ↗

(On Diff #455476)

can the vectors be shortened here? Does it also crash with a poison mask?

[Peter]

llvm/test/CodeGen/AArch64/aarch64-shuffle-undef.ll

24 ↗

(On Diff #455476)

From my experience, it seems only lengths of 32 will trigger the bug.

I just tested again. undef crashes both Debug and Release build. poison only crashes Debug build. This just adds the mystery for me...

I checked the failed testcases in Transform, it seems they only considered when part of the Mask is undef or poison, not the whole Mask.

[Peter]

llvm/test/CodeGen/AArch64/aarch64-shuffle-undef.ll

24 ↗

(On Diff #455476)

Ok I managed to find a testcase with only vector length of 16.

; ModuleID = 'PoC.ll'
source_filename = "M"

define void @f() {
BB:
  %A2 = alloca <16 x i32>, align 64
  %B = urem <16 x i32> <i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1>, <i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1>
  %S = shufflevector <16 x i32> %B, <16 x i32> %B, <16 x i32> undef
  store <16 x i32> %S, <16 x i32>* %A2, align 64
  ret void
}

What I don't understand is, if we change this case's integer type from i32 to i16, llc will complaint, even though the instruction seems legit to me

./build-debug/bin/llc: error: ./build-debug/bin/llc: PoC.ll:8:22: error: invalid shufflevector operands
  %S = shufflevector <16 x i16> %B, <16 x i16> %B, <16 x i16> undef

But when switching to <16 x i32>, it crashes...

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

This is probably too restrictive and not really matches the comment.

There are a few test failures with the patch, including LLVM.Transforms/InterleavedAccess/AArch64::interleaved-accesses.ll

llvm/test/CodeGen/AArch64/aarch64-shuffle-undef.ll

24 ↗

(On Diff #455476)

Could be a bit simplified further: https://llvm.godbolt.org/z/663P9z73G

Yeah, the vector of indices needs to have i32 as element type I think.

[Peter]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

I didn't know if undef or poison are part of the Mask is legal. That' s why I put a restrict check here.

If they are legal, should we just check the index every time before Mask is indexed? I am afraid that code won't be pretty, but that' s the only thing I can think of.

llvm/test/CodeGen/AArch64/aarch64-shuffle-undef.ll

24 ↗

(On Diff #455476)

That example is awesome. I though it needs to be a constant for it to work.

[fhahn]

Marking this as changes requested as there are some test failures.

[Peter]

Use new checks to sanitize Mask. Adding new tests and fixed test failures.

Previous check abort if anyone index is invalid, which is too strict.
A Mask can have undef/poison values in some of the indeces and we shouldn't abort in that case.
Thus we have failed testcases.

All we need to look out for is when the whole mask is undef, thus -1.
Therefore, we changed the checks from "abort if ANY index is invalid" to "abort if ALL index is invalid"

After discussion with Fhahn, we find that poison can also crash the code like undef, thus a poison testcase is added.

[Peter]

Fixed a formatting error

[Peter]

Fixed format issue

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

LLVM provides any_of/all_of helpers that take a iterator range. It would be good to use them here.

// If mask is `undef` or `poison`, `Mask` may be a vector of -1s.

There's a constexpr for UndefMaskElem: https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/Instructions.h#L1996

It would probably be better to check for that. IIUC this is the only case where an index could be out-of-range; other cases would be already rejected by the verifier I think.

[Peter]

Using llvm::all_of and UndefMaskElem to do the testing

[Peter]

Fixed format issue.

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

Don't we want any_of here? Won't a single -1 create an issue in the mask?

[Peter]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

No, It seems a few UndefMaskElem won't affect how it works. See test case https://github.com/llvm/llvm-project/blob/main/llvm/test/Transforms/InterleavedAccess/AArch64/interleaved-accesses-inseltpoison.ll.

Only a whole vector of UndefMaskElem will crash it.

[Peter]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

Correction, see test case here: https://github.com/llvm/llvm-project/blob/main/llvm/test/Transforms/InterleavedAccess/AArch64/interleaved-accesses-inseltpoison.ll#L246

Seems they expected a few elements to be undef or poison when they invented this feature, but never thought if the whole mask is undef.

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

Ah right, that makes me think that the patch doesn't fully fix the underlying issue.

E.g. the following below should also crash

define void @foo(<8 x i64> %a, ptr %dst) {
BB:
  %S = shufflevector <8 x i64> %a, <8 x i64> %a, <16 x i32> <i32 0, i32 undef, i32 undef, i32 undef,i32 undef, i32 undef, i32 undef, i32 undef,i32 undef, i32 undef, i32 undef, i32 undef,i32 undef, i32 undef, i32 undef, i32 undef>
  store <16 x i64> %S, ptr %dst, align 64
  ret void
}

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

If some lanes are undef, that's fine and there's already some support, the issue is that we compute an invalid index to find the start mask. But if all elements we need the select have undef indices, that's would still be fine.

There are more variations with only some undef/poison lanes that crash.

[Peter]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13369

You are right, the example you proposed is crashing.
Now I am lostdon't think I fully grasp the reason for the crashing.

[dmgreen]

This fails too: https://godbolt.org/z/f91z97cPY
The index it picks when the the first item in the mask is undef for a part of the split store.

[Peter]

I spent another morning trying to figure out how to fix this and failed.
I am confused why index is calculated as IdxJ * Factor + IdxI, isn't IdxJ = StoreCount * LaneLen * Factor + j; already multiplied by Factor? Why do it twice?
It seems to me there is no better way than inserting an

if (IdxJ * Factor + IdxI > Mask.size()) { 
  return true;
}

before every access.

If you agree with me I'll upload it.

[Peter]

Another question is, what is the expected outcome when the start index is undef?
It seems like when NEON is turned off, the whole shuffle will be ignored. Then we can check all start indices and abort if any of them is undef, but I am not sure if that's desired.

[dmgreen]

I am confused why index is calculated as IdxJ * Factor + IdxI, isn't IdxJ = StoreCount * LaneLen * Factor + j; already multiplied by Factor? Why do it twice?

Yeah - I think that logic is just wrong, so it is going out of range. Mask elements are almost never undef, and the combination of an undef first elements and wider-than-legal vectors is rare enough that no-one has noticed it is incorrect before. It would be good to support it though, if we can get the logic correct.

The Mask[IdxJ * Factor + IdxI] index into Mask should never be out-of-range. It needs to find the first element that isn't undef (-1) in that store-chunk, for elements in that Factor, and use that to find the equivalent StartMask. I think it should just be StoreCount * LaneLen * Factor + j * Factor + i, but I may have the LaneLen and Factor backwards. There is a comment that suggests that at least one of the elements shouldn't be undef, but it appears from your test if it is getting there then that might not be true.

[Peter]

fixed index calculation error.

[Peter]

In D132634#3825541, @dmgreen wrote:

I am confused why index is calculated as IdxJ * Factor + IdxI, isn't IdxJ = StoreCount * LaneLen * Factor + j; already multiplied by Factor? Why do it twice?

Yeah - I think that logic is just wrong, so it is going out of range. Mask elements are almost never undef, and the combination of an undef first elements and wider-than-legal vectors is rare enough that no-one has noticed it is incorrect before. It would be good to support it though, if we can get the logic correct.

The Mask[IdxJ * Factor + IdxI] index into Mask should never be out-of-range. It needs to find the first element that isn't undef (-1) in that store-chunk, for elements in that Factor, and use that to find the equivalent StartMask. I think it should just be StoreCount * LaneLen * Factor + j * Factor + i, but I may have the LaneLen and Factor backwards. There is a comment that suggests that at least one of the elements shouldn't be undef, but it appears from your test if it is getting there then that might not be true.

I am more confident to fix it once you feel like it's a bug too. I just submitted latest change.

[dmgreen]

Thanks - this looks OK to me. Can you add these two extra tests to the new test file, both of which will produce the same output. Then this LGTM.

define void @noundefs(<8 x i32> %a, <8 x i32> %b, ptr %dst) {
BB:
  %S = shufflevector <8 x i32> %a, <8 x i32> %b, <16 x i32> <i32 0, i32 8, i32 1, i32 9, i32 2, i32 10, i32 3, i32 11, i32 4, i32 12, i32 5, i32 13, i32 undef, i32 14, i32 7, i32 15>
  store <16 x i32> %S, ptr %dst, align 64
  ret void
}

define void @undefs(<8 x i32> %a, <8 x i32> %b, ptr %dst) {
BB:
  %S = shufflevector <8 x i32> %a, <8 x i32> %b, <16 x i32> <i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 3, i32 11, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 7, i32 15>
  store <16 x i32> %S, ptr %dst, align 64
  ret void
}

llvm/test/CodeGen/AArch64/aarch64-shufflevector.ll

1

aarch64-shufflevector.ll is a bit of a generic name for this test file, considering what it is testing. Maybe change it to be something that mentions interleave stores with undef elements?

[Peter]

add more tests. change file name to indicate interleaved access w under

[dmgreen]

Thanks. LGTM

[Peter]

@fhahn Hi Florian, when you have time can you review this and approve it when you have time?

I don't have push access to LLVM upstream, I'm hoping if you can push the diff for me or grant me access.

[dmgreen]

I can commit it, so long as you can provide a "name <[e-mail]>" to use for attribution of the git commit.

[Peter]

If you can do it that's the best.
Please use "Peter Rong, [e-mail]".

Thanks.

[fhahn]

Thanks for the latest changes, LGTM!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (NumStores > 1) {
     // element type.
     BaseAddr = Builder.CreateBitCast(
         BaseAddr,
         SubVecTy->getElementType()->getPointerTo(SI->getPointerAddressSpace()));
   }
 
   auto Mask = SVI->getShuffleMask();
 
+  // Sanity check if all the indices are NOT in range.
+  // If mask is `undef` or `poison`, `Mask` may be a vector of -1s.
+  // If all of them are `undef`, OOB read will happen later.
+  if (llvm::all_of(Mask, [](int Idx) { return Idx == UndefMaskElem; })) {
+    return false;
+  }
+
   Type *PtrTy =
       UseScalable
           ? STVTy->getElementType()->getPointerTo(SI->getPointerAddressSpace())
           : STVTy->getPointerTo(SI->getPointerAddressSpace());
   Type *PredTy = VectorType::get(Type::getInt1Ty(STVTy->getContext()),
                                  STVTy->getElementCount());
 
   static const Intrinsic::ID SVEStoreIntrs[3] = {Intrinsic::aarch64_sve_st2,
for (unsigned i = 0; i < Factor; i++) {
       Value *Shuffle;
       unsigned IdxI = StoreCount * LaneLen * Factor + i;
       if (Mask[IdxI] >= 0) {
         Shuffle = Builder.CreateShuffleVector(
             Op0, Op1, createSequentialMask(Mask[IdxI], LaneLen, 0));
       } else {
         unsigned StartMask = 0;
         for (unsigned j = 1; j < LaneLen; j++) {
-          unsigned IdxJ = StoreCount * LaneLen * Factor + j;
-          if (Mask[IdxJ * Factor + IdxI] >= 0) {
-            StartMask = Mask[IdxJ * Factor + IdxI] - IdxJ;
+          unsigned IdxJ = StoreCount * LaneLen * Factor + j * Factor + i;
+          if (Mask[IdxJ] >= 0) {
+            StartMask = Mask[IdxJ] - j;
             break;
           }
         }
         // 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

llvm/test/CodeGen/AArch64/aarch64-shufflevector.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -opaque-pointers < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+define void @f_undef(<8 x i64> %a, ptr %dst) {
+; CHECK-LABEL: f_undef:
+; CHECK:       // %bb.0: // %BB
+; CHECK-NEXT:    ret
+BB:
+  %S = shufflevector <8 x i64> %a, <8 x i64> %a, <16 x i32> undef
+  store <16 x i64> %S, ptr %dst, align 64
+  ret void
+}
+
+define void @f_poison(<8 x i64> %a, ptr %dst) {
+; CHECK-LABEL: f_poison:
+; CHECK:       // %bb.0: // %BB
+; CHECK-NEXT:    ret
+BB:
+  %S = shufflevector <8 x i64> %a, <8 x i64> %a, <16 x i32> poison
+  store <16 x i64> %S, ptr %dst, align 64
+  ret void
+}
+
+define void @f_undef_15(<8 x i64> %a, ptr %dst) {
+; CHECK-LABEL: f_undef_15:
+; CHECK:       // %bb.0: // %BB
+; CHECK-NEXT:    mov x9, x0
+; CHECK-NEXT:    // kill: def $q0 killed $q0 def $q0_q1
+; CHECK-NEXT:    add x8, x0, #64
+; CHECK-NEXT:    mov v1.16b, v0.16b
+; CHECK-NEXT:    st2 { v0.2d, v1.2d }, [x9], #32
+; CHECK-NEXT:    st2 { v0.2d, v1.2d }, [x8]
+; CHECK-NEXT:    add x8, x0, #96
+; CHECK-NEXT:    st2 { v0.2d, v1.2d }, [x9]
+; CHECK-NEXT:    st2 { v0.2d, v1.2d }, [x8]
+; CHECK-NEXT:    ret
+BB:
+  %S = shufflevector <8 x i64> %a, <8 x i64> %a, <16 x i32> <i32 0, i32 undef, i32 undef, i32 undef,i32 undef, i32 undef, i32 undef, i32 undef,i32 undef, i32 undef, i32 undef, i32 undef,i32 undef, i32 undef, i32 undef, i32 undef>
+  store <16 x i64> %S, ptr %dst, align 64
+  ret void
+}
+
+define void @f_undef_1(<8 x i64> %a, ptr %dst) {
+; CHECK-LABEL: f_undef_1:
+; CHECK:       // %bb.0: // %BB
+; CHECK-NEXT:    mov x9, x0
+; CHECK-NEXT:    add x8, x0, #64
+; CHECK-NEXT:    mov v16.16b, v0.16b
+; CHECK-NEXT:    // kill: def $q3 killed $q3 def $q3_q4
+; CHECK-NEXT:    mov v17.16b, v16.16b
+; CHECK-NEXT:    mov v5.16b, v2.16b
+; CHECK-NEXT:    // kill: def $q1 killed $q1 def $q1_q2
+; CHECK-NEXT:    st2 { v16.2d, v17.2d }, [x9], #32
+; CHECK-NEXT:    mov v6.16b, v5.16b
+; CHECK-NEXT:    mov v2.16b, v1.16b
+; CHECK-NEXT:    mov v4.16b, v3.16b
+; CHECK-NEXT:    st2 { v5.2d, v6.2d }, [x8]
+; CHECK-NEXT:    add x8, x0, #96
+; CHECK-NEXT:    st2 { v1.2d, v2.2d }, [x9]
+; CHECK-NEXT:    st2 { v3.2d, v4.2d }, [x8]
+; CHECK-NEXT:    ret
+BB:
+  %S = shufflevector <8 x i64> %a, <8 x i64> %a, <16 x i32> <i32 0, i32 8, i32 1, i32 9, i32 2, i32 10, i32 3, i32 11, i32 4, i32 12, i32 5, i32 13, i32 undef, i32 14, i32 7, i32 15>
+  store <16 x i64> %S, ptr %dst, align 64
+  ret void
+}