[SVE] Fix TypeSize->uint64_t implicit conversion in visitAlloca()

[kmclaughlin]

Summary

Fixes a crash ('Invalid size request on a scalable vector') in visitAlloca()
when we call this function for a scalable alloca instruction, caused
by the implicit conversion of TySize to uint64_t.
This patch changes TySize to a TypeSize as returned by getTypeAllocSize()
and ensures the allocation size is multiplied by vscale for scalable vectors.

[david-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

4028

Hi @kmclaughlin, I think the known min value needs multiplying by vscale here. You can use DAG.getVScale() to do this I think.

llvm/test/CodeGen/AArch64/sve-alloca.ll

6

It would be good to add some check lines for the stack allocation here too if possible?

[madhur13490]

Can we please have a bit more detailed description for this patch? What problem this patch is trying to solve?

[kmclaughlin]

• Fixed AllocSize to ensure we multiply by vscale for scalable vectors
• Added CHECK lines to the test in sve-alloca.ll

[sdesmalen]

llvm/test/CodeGen/AArch64/sve-alloca.ll

1

Hi @kmclaughlin, it seems there are some CHECK lines missing below, because the check for the condition has disappeared. Please don't remove any CHECK lines if you keep:

; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py

in the top of the file.

i.e. either don't keep this line, or don't manually edit the tests after generating.

In this case, I prefer to keep all context with all CHECK lines.

[sdesmalen]

LGTM

llvm/test/CodeGen/AArch64/sve-alloca.ll

10–14

I think the explicit alignment could be folded away if LLVM would know something about vscale from the vscale_range attribute (through computeKnownBits?), although we'd need to know if vscale is a power of 2, which isn't encoded by vscale_range unless min(vscale_range) == max(vscale_range). In any case, because VSCALE is currently an opaque node, we always end up with the alignment code.

[kmclaughlin]

• Generated the check lines in sve-alloca.ll with update_llc_test_checks.py
• Use TySize.getFixedValue() when TySize is not scalable

[david-arm]

LGTM!

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

void SelectionDAGBuilder::visitAlloca(const AllocaInst &I) {
   // allocate it statically on the stack.
   if (FuncInfo.StaticAllocaMap.count(&I))
     return;   // getValue will auto-populate this.
 
   SDLoc dl = getCurSDLoc();
   Type *Ty = I.getAllocatedType();
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   auto &DL = DAG.getDataLayout();
-  uint64_t TySize = DL.getTypeAllocSize(Ty);
+  TypeSize TySize = DL.getTypeAllocSize(Ty);
   MaybeAlign Alignment = std::max(DL.getPrefTypeAlign(Ty), I.getAlign());
 
   SDValue AllocSize = getValue(I.getArraySize());
 
   EVT IntPtr = TLI.getPointerTy(DAG.getDataLayout(), DL.getAllocaAddrSpace());
   if (AllocSize.getValueType() != IntPtr)
     AllocSize = DAG.getZExtOrTrunc(AllocSize, dl, IntPtr);
 
-  AllocSize = DAG.getNode(ISD::MUL, dl, IntPtr,
-                          AllocSize,
+  if (TySize.isScalable())
+    AllocSize = DAG.getNode(ISD::MUL, dl, IntPtr, AllocSize,
+                            DAG.getVScale(dl, IntPtr,
+                                          APInt(IntPtr.getScalarSizeInBits(),
+                                                TySize.getKnownMinValue())));
+  else
+    AllocSize = DAG.getNode(ISD::MUL, dl, IntPtr, AllocSize,
                             DAG.getConstant(TySize, dl, IntPtr));
 
   // Handle alignment.  If the requested alignment is less than or equal to
   // the stack alignment, ignore it.  If the size is greater than or equal to
   // the stack alignment, we note this in the DYNAMIC_STACKALLOC node.
   Align StackAlign = DAG.getSubtarget().getFrameLowering()->getStackAlign();
   if (*Alignment <= StackAlign)
     Alignment = None;
 

llvm/test/CodeGen/AArch64/sve-alloca.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64 -mattr=+sve < %s | FileCheck %s
+
+declare void @bar(<vscale x 4 x i64>*)
+
+define void @foo(<vscale x 4 x i64> %dst, i1 %cond) {
+; CHECK-LABEL: foo:
+; CHECK:         rdvl x9, #2
+; CHECK-NEXT:    mov x8, sp
+; CHECK-NEXT:    add x9, x9, #15
+; CHECK-NEXT:    and x9, x9, #0xfffffffffffffff0
+; CHECK-NEXT:    sub x8, x8, x9
+; CHECK-NEXT:    and x0, x8, #0xffffffffffffffe0
+; CHECK-NEXT:    mov sp, x0
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    st1d { z1.d }, p0, [x0, #1, mul vl]
+; CHECK-NEXT:    st1d { z0.d }, p0, [x0]
+; CHECK-NEXT:    bl bar
+entry:
+  br i1 %cond, label %if.then, label %if.end
+
+if.then:
+  %ptr = alloca <vscale x 4 x i64>
+  store <vscale x 4 x i64> %dst, <vscale x 4 x i64>* %ptr
+  call void @bar(<vscale x 4 x i64>* %ptr)
+  br label %if.end
+
+if.end:
+  ret void
+}