[RISCV] Scale scalably-typed split argument offsets by VSCALE

[frasercrmck]

Summary

This patch fixes a bug in lowering scalable-vector types in RISC-V's
main calling convention. When scalable-vector types are split and passed
indirectly, the target is responsible for scaling the offset --
initially set to the known-minimum store size -- by the scalable factor.

Before this we were issuing overlapping loads or stores to the different
parts, leading to incorrect codegen.

Credit to @HsiangKai for spotting this.

[frasercrmck]

• remove FIXMEs, on account of it being.. fixed

[frasercrmck]

llvm/test/CodeGen/RISCV/rvv/calling-conv.ll

50

I'd better dig into these 32s (presumably the "loc mem offset"?) and make sure they're going to "scale" to multiple split arguments.

[frasercrmck]

llvm/test/CodeGen/RISCV/rvv/calling-conv.ll

50

Looks like they're fine as they're regular stack FIs which are correctly typed as scalable-vector.

define <vscale x 32 x i32> @caller(<vscale x 32 x i32> %x) {
; RV64-NEXT:    ret
  %a = call <vscale x 32 x i32> @callee(<vscale x 32 x i32> zeroinitializer, <vscale x 32 x i32> %x, <vscale x 32 x i32> zeroinitializer)

stack:
  - { id: 0, name: '', type: default, offset: 0, size: 128, alignment: 128,
      stack-id: scalable-vector, callee-saved-register: '', callee-saved-restored: true,
      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
  - { id: 1, name: '', type: default, offset: 0, size: 128, alignment: 128,
      stack-id: scalable-vector, callee-saved-register: '', callee-saved-restored: true,
      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }

...
  bb.0 (%ir-block.0):
    liveins: $v8m8, $v16m8

    %1:vrm8 = COPY $v16m8
    %0:vrm8 = COPY $v8m8
    %2:gpr = PseudoReadVLENB
    %3:gpr = SLLI killed %2, 3
    %4:gpr = ADDI %stack.1, 0
    %5:gpr = ADD %4, %3
    ADJCALLSTACKDOWN 0, 0, implicit-def dead $x2, implicit $x2
    %6:vrm8 = PseudoVMV_V_I_M8 0, $x0, 5
    VS8R_V %6, killed %5 :: (store unknown-size into %stack.1, align 64)
    %7:gpr = ADDI %stack.0, 0
    %8:gpr = ADD %7, %3
    VS8R_V %1, killed %8 :: (store unknown-size into %stack.0, align 64)
    VS8R_V %6, %stack.1 :: (store unknown-size into %stack.1, align 128)
    VS8R_V %0, %stack.0 :: (store unknown-size into %stack.0, align 128)
    $v8m8 = COPY %6
    $v16m8 = COPY %6
    $x10 = COPY %7
    $x12 = COPY %4
    PseudoCALL target-flags(riscv-plt) @callee, csr_ilp32_lp64, implicit-def dead $x1, implicit $v8m8, implicit $v16m8, implicit $x10, implicit $x12, implicit-def $x2, implicit-def $v8m8, implicit-def $v16m8

[HsiangKai]

LGTM.

[frasercrmck]

• rebase

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

if (VA.getLocInfo() == CCValAssign::Indirect) {
       InVals.push_back(DAG.getLoad(VA.getValVT(), DL, Chain, ArgValue,
                                    MachinePointerInfo()));
       unsigned ArgIndex = Ins[i].OrigArgIndex;
       unsigned ArgPartOffset = Ins[i].PartOffset;
       assert(VA.getValVT().isVector() || ArgPartOffset == 0);
       while (i + 1 != e && Ins[i + 1].OrigArgIndex == ArgIndex) {
         CCValAssign &PartVA = ArgLocs[i + 1];
         unsigned PartOffset = Ins[i + 1].PartOffset - ArgPartOffset;
-        SDValue Address = DAG.getNode(ISD::ADD, DL, PtrVT, ArgValue,
-                                      DAG.getIntPtrConstant(PartOffset, DL));
+        SDValue Offset = DAG.getIntPtrConstant(PartOffset, DL);
+        if (PartVA.getValVT().isScalableVector())
+          Offset = DAG.getNode(ISD::VSCALE, DL, XLenVT, Offset);
+        SDValue Address = DAG.getNode(ISD::ADD, DL, PtrVT, ArgValue, Offset);
         InVals.push_back(DAG.getLoad(PartVA.getValVT(), DL, Chain, Address,
                                      MachinePointerInfo()));
         ++i;
       }
       continue;
     }
     InVals.push_back(ArgValue);
   }
if (VA.getLocInfo() == CCValAssign::Indirect) {
       // which case the base address is partly offset and subsequent stores are
       // relative to that.
       unsigned ArgIndex = Outs[i].OrigArgIndex;
       unsigned ArgPartOffset = Outs[i].PartOffset;
       assert(VA.getValVT().isVector() || ArgPartOffset == 0);
       // Calculate the total size to store. We don't have access to what we're
       // actually storing other than performing the loop and collecting the
       // info.
-      SmallVector<std::pair<SDValue, unsigned>> Parts;
+      SmallVector<std::pair<SDValue, SDValue>> Parts;
       while (i + 1 != e && Outs[i + 1].OrigArgIndex == ArgIndex) {
         SDValue PartValue = OutVals[i + 1];
         unsigned PartOffset = Outs[i + 1].PartOffset - ArgPartOffset;
+        SDValue Offset = DAG.getIntPtrConstant(PartOffset, DL);
         EVT PartVT = PartValue.getValueType();
+        if (PartVT.isScalableVector())
+          Offset = DAG.getNode(ISD::VSCALE, DL, XLenVT, Offset);
         StoredSize += PartVT.getStoreSize();
         StackAlign = std::max(StackAlign, getPrefTypeAlign(PartVT, DAG));
-        Parts.push_back(std::make_pair(PartValue, PartOffset));
+        Parts.push_back(std::make_pair(PartValue, Offset));
         ++i;
       }
       SDValue SpillSlot = DAG.CreateStackTemporary(StoredSize, StackAlign);
       int FI = cast<FrameIndexSDNode>(SpillSlot)->getIndex();
       MemOpChains.push_back(
           DAG.getStore(Chain, DL, ArgValue, SpillSlot,
                        MachinePointerInfo::getFixedStack(MF, FI)));
       for (const auto &Part : Parts) {
         SDValue PartValue = Part.first;
-        unsigned PartOffset = Part.second;
-        SDValue Address = DAG.getNode(ISD::ADD, DL, PtrVT, SpillSlot,
-                                      DAG.getIntPtrConstant(PartOffset, DL));
+        SDValue PartOffset = Part.second;
+        SDValue Address =
+            DAG.getNode(ISD::ADD, DL, PtrVT, SpillSlot, PartOffset);
         MemOpChains.push_back(
             DAG.getStore(Chain, DL, PartValue, Address,
                          MachinePointerInfo::getFixedStack(MF, FI)));
       }
       ArgValue = SpillSlot;
     } else {
       ArgValue = convertValVTToLocVT(DAG, ArgValue, VA, DL, Subtarget);
     }

llvm/test/CodeGen/RISCV/rvv/calling-conv.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -mattr=+m,+experimental-v < %s | FileCheck %s --check-prefix=RV32
 ; RUN: llc -mtriple=riscv64 -mattr=+m,+experimental-v < %s | FileCheck %s --check-prefix=RV64
 
 ; Check that we correctly scale the split part indirect offsets by VSCALE.
-; FIXME: We don't; we're loading a full 8 vector registers 64 bytes ahead of
-; the first address. This should be scaled by vlenb!
 define <vscale x 32 x i32> @callee_scalable_vector_split_indirect(<vscale x 32 x i32> %x, <vscale x 32 x i32> %y) {
 ; RV32-LABEL: callee_scalable_vector_split_indirect:
 ; RV32:       # %bb.0:
-; RV32-NEXT:    addi a1, a0, 64
+; RV32-NEXT:    csrr a1, vlenb
+; RV32-NEXT:    slli a1, a1, 3
+; RV32-NEXT:    add a1, a0, a1
 ; RV32-NEXT:    vl8re32.v v24, (a0)
 ; RV32-NEXT:    vl8re32.v v0, (a1)
 ; RV32-NEXT:    vsetvli a0, zero, e32,m8,ta,mu
 ; RV32-NEXT:    vadd.vv v8, v8, v24
 ; RV32-NEXT:    vadd.vv v16, v16, v0
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: callee_scalable_vector_split_indirect:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    addi a1, a0, 64
+; RV64-NEXT:    csrr a1, vlenb
+; RV64-NEXT:    slli a1, a1, 3
+; RV64-NEXT:    add a1, a0, a1
 ; RV64-NEXT:    vl8re32.v v24, (a0)
 ; RV64-NEXT:    vl8re32.v v0, (a1)
 ; RV64-NEXT:    vsetvli a0, zero, e32,m8,ta,mu
 ; RV64-NEXT:    vadd.vv v8, v8, v24
 ; RV64-NEXT:    vadd.vv v16, v16, v0
 ; RV64-NEXT:    ret
   %a = add <vscale x 32 x i32> %x, %y
   ret <vscale x 32 x i32> %a
 }
 
 ; Call the function above. Check that we set the arguments correctly.
-; FIXME: We don't, see above.
 define <vscale x 32 x i32> @caller_scalable_vector_split_indirect(<vscale x 32 x i32> %x) {
 ; RV32-LABEL: caller_scalable_vector_split_indirect:
 ; RV32:       # %bb.0:
 ; RV32-NEXT:    addi sp, sp, -48
 ; RV32-NEXT:    .cfi_def_cfa_offset 48
 ; RV32-NEXT:    sw ra, 44(sp) # 4-byte Folded Spill
 ; RV32-NEXT:    .cfi_offset ra, -4
 ; RV32-NEXT:    csrr a0, vlenb
 ; RV32-NEXT:    slli a0, a0, 4
 ; RV32-NEXT:    sub sp, sp, a0
-; RV32-NEXT:    addi a0, sp, 96
+; RV32-NEXT:    csrr a0, vlenb
+; RV32-NEXT:    slli a0, a0, 3
+; RV32-NEXT:    addi a1, sp, 32
+; RV32-NEXT:    add a0, a1, a0
 ; RV32-NEXT:    vs8r.v v16, (a0)
 ; RV32-NEXT:    addi a0, sp, 32
 ; RV32-NEXT:    vs8r.v v8, (a0)
 ; RV32-NEXT:    vsetvli a0, zero, e32,m8,ta,mu
 ; RV32-NEXT:    vmv.v.i v8, 0
 ; RV32-NEXT:    addi a0, sp, 32
 ; RV32-NEXT:    vmv8r.v v16, v8
 ; RV32-NEXT:    call callee_scalable_vector_split_indirect@plt
 ; RV32-NEXT:    csrr a0, vlenb
 ; RV32-NEXT:    slli a0, a0, 4
 ; RV32-NEXT:    add sp, sp, a0
 ; RV32-NEXT:    lw ra, 44(sp) # 4-byte Folded Reload
 ; RV32-NEXT:    addi sp, sp, 48
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: caller_scalable_vector_split_indirect:
 ; RV64:       # %bb.0:
 ; RV64-NEXT:    addi sp, sp, -32
 ; RV64-NEXT:    .cfi_def_cfa_offset 32
 ; RV64-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
 ; RV64-NEXT:    .cfi_offset ra, -8
 ; RV64-NEXT:    csrr a0, vlenb
 ; RV64-NEXT:    slli a0, a0, 4
 ; RV64-NEXT:    sub sp, sp, a0
-; RV64-NEXT:    addi a0, sp, 88
+; RV64-NEXT:    csrr a0, vlenb
+; RV64-NEXT:    slli a0, a0, 3
+; RV64-NEXT:    addi a1, sp, 24
+; RV64-NEXT:    add a0, a1, a0
 ; RV64-NEXT:    vs8r.v v16, (a0)
 ; RV64-NEXT:    addi a0, sp, 24
 ; RV64-NEXT:    vs8r.v v8, (a0)
 ; RV64-NEXT:    vsetvli a0, zero, e32,m8,ta,mu
 ; RV64-NEXT:    vmv.v.i v8, 0
 ; RV64-NEXT:    addi a0, sp, 24
 ; RV64-NEXT:    vmv8r.v v16, v8
 ; RV64-NEXT:    call callee_scalable_vector_split_indirect@plt