[SelectionDAG] Fix getVectorSubVecPointer for scalable subvectors.

[sdesmalen]

Summary

When inserting a scalable subvector into a scalable vector through
the stack, the index to store to needs to be scaled by vscale.
Before this patch, that didn't yet happen, so it would generate the
wrong offset, thus storing a subvector to the incorrect address
and overwriting the wrong lanes.

For some insert:

nxv8f16 insert_subvector(nxv8f16 %vec, nxv2f16 %subvec, i64 2)

The offset was not scaled by vscale:

orr     x8, x8, #0x4
st1h    { z0.h }, p0, [sp]
st1h    { z1.d }, p1, [x8]
ld1h    { z0.h }, p0/z, [sp]

And is changed to:

mov x8, sp
st1h { z0.h }, p0, [sp]
st1h { z1.d }, p1, [x8, #1, mul vl]
ld1h { z0.h }, p0/z, [sp]

[CarolineConcatto]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7874

Can you do something like this:

Index = DAG.getVScale(DL, IdxVT, APInt(IdxVT.getSizeInBits(), Index.getConstantOperandVal(0)));

?
Insead of multiply the index by a scalable vector of size 1?

[david-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7867–7868

Hi @sdesmalen, it's just a thought, but while you're in this area is it also worth clamping the index for scalable vectors too? The comment above is incorrect, because we do explicitly clamp the index in other places for scalable vectors.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7867–7868

I don't think any clamping is required, because when both the subvector and the vector being inserted into are scalable, we know at compiletime whether the vector index will exceed the size of the input vector.

[david-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7867–7868

Oh ok - I wonder why we do this for fixed-length vectors? I was sort of expecting the problem to be the same for both inserting fixed into fixed and inserting scalable into scalable?

I was specifically worried about what we did in practice for this case:

call <vscale x 8 x half> @llvm.experimental.vector.insert.nxv8f16.nxv2f16(<vscale x 8 x half> %vec, <vscale x 2 x half> %in, i64 10)

because if vscale=1 then we're inserting beyond the end of the vector.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7867–7868

This exists for inserting/extracting fixed from scalable, where we don't know at compile-time if the fixed offset exceeds the scalable vector.

The example you give here:

call <vscale x 8 x half> @llvm.experimental.vector.insert.nxv8f16.nxv2f16(<vscale x 8 x half> %vec, <vscale x 2 x half> %in, i64 10)

is always out of bounds, because vscale*10 > vscale*(8-2) for any vscale.

[david-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7867–7868

OK, I still wasn't sure what you meant here by checking at compile time, but I tried out a test manually and I see that the Verifier emits an error for indices that are too large. So that's fine then!

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7874

I'm not sure if you can always know for sure that Index is a constant value so that's why I used an explicit multiply with vscale.

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7874

To me the rational for not using getConstantOperandAPInt is at odds to the rational for not clamping the index. None of what you say is wrong but there's nothing to say this function has to be used in conjunction with ISD::INSERT_SUBVECTOR, but the function's description does say:

If \p Idx plus the size of \p SubVecVT is out of bounds the returned pointer is unspecified, but the value returned will be such that the entire subvector would be within the vector bounds.

So either this function is only ever used in combination with EXTRACT_SUBVECTOR/INSERT_SUBVECTOR, in which case we can assume Index to be a constant (perhaps even changing the prototype to force this?), or this is a generic helper function and thus Index can be anything and must be clamped to honour the function's description.

[sdesmalen]

• Simplified getVectorSubVecPointer to always clamp dynamically.
• For the combinations in the table below, this revision changed clampDynamicVectorIndex as follows:
  • Accepts cases: A, B, C, D, E, F (C and F were added)
  • It now also asserts that A and F must have valid indices (i.e. are already validated by their original operation, e.g. EXTRACT/INSERT_SUBVECTOR or EXTRACT/INSERT_VECTOR_ELT.

	  |  index  |  subvector  |  vector
	  |----------------------------------
	A |  const        fixed         fixed
	B |  const        fixed      scalable
	C |  const     scalable      scalable
	D |    var        fixed         fixed
	E |    var        fixed      scalable
	F |    var     scalable      scalable
	
	Note that the following combinations are invalid, because we don't support
	extracting a scalable vector from a fixed-width vector:
	X |  const     scalable         fixed
	X |    var     scalable         fixed

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7822–7824

Given you're not making assumptions as to where Idx is coming from I don't think an assert is safe enough. Sure an asserts build will exit here but a release build could leak/corrupt data, which is a problem the user is trying to prevent, hence calling this function.

Instead I think the assert should be replaced by always clamping Idx (see the MaxIndex calculation below). If the source of Idx is indeed an EXTRACT_SUBVECTOR/INSERT_SUBVECTOR then the invalid index should be asserting as part of getNode rather than getting this far.

[sdesmalen]

Always clamp instead of asserting.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7822–7824

Good point about the generated code possibly doing the wrong thing for non-assert builds. I've changed it to always clamp, and removed the assert. From looking at the uses of getVectorSubVecPointer, it's not really worth passing in some bool to tell whether the value is coming from an insert/extract subvector in order to assert. The nodes where this function is called for insert/extract subvectors, have nodes that must already have been checked for correctness in other places.

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7869

The assert text is no longer correct.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

7869

Good spot, I've changed it! Thanks.

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

TargetLowering::IncrementMemoryAddress(SDValue Addr, SDValue Mask,
   } else
     Increment = DAG.getConstant(DataVT.getStoreSize(), DL, AddrVT);
 
   return DAG.getNode(ISD::ADD, DL, AddrVT, Addr, Increment);
 }
 
 static SDValue clampDynamicVectorIndex(SelectionDAG &DAG, SDValue Idx,
                                        EVT VecVT, const SDLoc &dl,
-                                       unsigned NumSubElts) {
-  if (!VecVT.isScalableVector() && isa<ConstantSDNode>(Idx))
-    return Idx;
+                                       ElementCount SubEC) {
+  assert(!(SubEC.isScalable() && VecVT.isFixedLengthVector()) &&
+         "Cannot index a scalable vector within a fixed-width vector");
 
-  EVT IdxVT = Idx.getValueType();
   unsigned NElts = VecVT.getVectorMinNumElements();
-  if (VecVT.isScalableVector()) {
+  unsigned NumSubElts = SubEC.getKnownMinValue();
+  EVT IdxVT = Idx.getValueType();
+
+  if (VecVT.isScalableVector() && !SubEC.isScalable()) {
     // If this is a constant index and we know the value plus the number of the
     // elements in the subvector minus one is less than the minimum number of
     // elements then it's safe to return Idx.
     if (auto *IdxCst = dyn_cast<ConstantSDNode>(Idx))
       if (IdxCst->getZExtValue() + (NumSubElts - 1) < NElts)
         return Idx;
     SDValue VS =
         DAG.getVScale(dl, IdxVT, APInt(IdxVT.getFixedSizeInBits(), NElts));
     unsigned SubOpcode = NumSubElts <= NElts ? ISD::SUB : ISD::USUBSAT;
     SDValue Sub = DAG.getNode(SubOpcode, dl, IdxVT, VS,
                               DAG.getConstant(NumSubElts, dl, IdxVT));
SDValue TargetLowering::getVectorSubVecPointer(SelectionDAG &DAG,
   // Make sure the index type is big enough to compute in.
   Index = DAG.getZExtOrTrunc(Index, dl, VecPtr.getValueType());
 
   EVT EltVT = VecVT.getVectorElementType();
 
   // Calculate the element offset and add it to the pointer.
   unsigned EltSize = EltVT.getFixedSizeInBits() / 8; // FIXME: should be ABI size.
   assert(EltSize * 8 == EltVT.getFixedSizeInBits() &&
          "Converting bits to bytes lost precision");
-
-  // Scalable vectors don't need clamping as these are checked at compile time
-  if (SubVecVT.isFixedLengthVector()) {
-    assert(SubVecVT.getVectorElementType() == EltVT &&
-           "Sub-vector must be a fixed vector with matching element type");
-    Index = clampDynamicVectorIndex(DAG, Index, VecVT, dl,
-                                    SubVecVT.getVectorNumElements());
-  }
+         "Sub-vector must be a vector with matching element type");
+                                  SubVecVT.getVectorElementCount());
 
   EVT IdxVT = Index.getValueType();
+  if (SubVecVT.isScalableVector())
+    Index =
+        DAG.getNode(ISD::MUL, dl, IdxVT, Index,
+                    DAG.getVScale(dl, IdxVT, APInt(IdxVT.getSizeInBits(), 1)));
 
   Index = DAG.getNode(ISD::MUL, dl, IdxVT, Index,
                       DAG.getConstant(EltSize, dl, IdxVT));
   return DAG.getMemBasePlusOffset(VecPtr, Index, dl);
 }
 
 //===----------------------------------------------------------------------===//
 // Implementation of Emulated TLS Model

llvm/test/CodeGen/AArch64/sve-insert-vector.ll

 ; CHECK-NEXT:    uunpklo z2.d, z2.s
 ; CHECK-NEXT:    uzp1 z1.s, z2.s, z1.s
 ; CHECK-NEXT:    uzp1 z0.h, z1.h, z0.h
 ; CHECK-NEXT:    ret
   %r = call <vscale x 8 x i16> @llvm.experimental.vector.insert.nxv8i16.nxv2i16(<vscale x 8 x i16> %vec, <vscale x 2 x i16> %in, i64 2)
   ret <vscale x 8 x i16> %r
 }
 
+; Test that the index is scaled by vscale if the subvector is scalable.
+define <vscale x 8 x half> @insert_nxv8f16_nxv2f16(<vscale x 8 x half> %vec, <vscale x 2 x half> %in) nounwind {
+; CHECK-LABEL: insert_nxv8f16_nxv2f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    addvl sp, sp, #-1
+; CHECK-NEXT:    mov x8, sp
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    st1h { z0.h }, p0, [sp]
+; CHECK-NEXT:    st1h { z1.d }, p1, [x8, #1, mul vl]
+; CHECK-NEXT:    ld1h { z0.h }, p0/z, [sp]
+; CHECK-NEXT:    addvl sp, sp, #1
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+  %r = call <vscale x 8 x half> @llvm.experimental.vector.insert.nxv8f16.nxv2f16(<vscale x 8 x half> %vec, <vscale x 2 x half> %in, i64 2)
+  ret <vscale x 8 x half> %r
+}
 
 ; Fixed length clamping
 
 define <vscale x 2 x i64> @insert_fixed_v2i64_nxv2i64(<vscale x 2 x i64> %vec, <2 x i64> %subvec) nounwind #0 {
 ; CHECK-LABEL: insert_fixed_v2i64_nxv2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1

 declare <vscale x 2 x i64> @llvm.experimental.vector.insert.nxv2i64.v4i64(<vscale x 2 x i64>, <4 x i64>, i64)
 
 declare <vscale x 16 x i64> @llvm.experimental.vector.insert.nxv8i64.nxv16i64(<vscale x 16 x i64>, <vscale x 8 x i64>, i64)
 declare <vscale x 16 x i64> @llvm.experimental.vector.insert.v2i64.nxv16i64(<vscale x 16 x i64>, <2 x i64>, i64)
 declare <vscale x 4 x i32> @llvm.experimental.vector.insert.nxv4i32.nxv1i32(<vscale x 4 x i32>, <vscale x 1 x i32>, i64)
 declare <vscale x 6 x i16> @llvm.experimental.vector.insert.nxv6i16.nxv1i16(<vscale x 6 x i16>, <vscale x 1 x i16>, i64)
 
 declare <vscale x 8 x i16> @llvm.experimental.vector.insert.nxv8i16.nxv2i16(<vscale x 8 x i16>, <vscale x 2 x i16>, i64)
+
+declare <vscale x 8 x half> @llvm.experimental.vector.insert.nxv8f16.nxv2f16(<vscale x 8 x half>, <vscale x 2 x half>, i64)